import os
from .sensevoice_bin import SenseVoiceSmall
from .postprocess_utils import rich_transcription_postprocess


cur_dir = os.path.dirname(os.path.abspath(__file__))
asr_model_path = os.path.join(cur_dir, 'models')

class AsrModel:
    def __init__(
        self,
        model_dir=asr_model_path
    ):
        self._model = SenseVoiceSmall(
            model_dir,
            batch_size=10
        )

    def __call__(
        self, wav_file_path
    ):
        res = self._model(
            wav_file_path,
            language='zh',
            use_itn=True
        )
        res = res[0][0].tolist()
        text = rich_transcription_postprocess(res[0])
        return text